#include "rviz_custom_plugins/utils/bag_map_loader.h"
#include <rosbag2_cpp/reader.hpp>
#include <rclcpp/serialization.hpp>
#include <rclcpp/rclcpp.hpp>
#include <rviz_common/load_resource.hpp>
#include <tf2/LinearMath/Transform.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <geometry_msgs/msg/pose.hpp>
#include <tf2_sensor_msgs/tf2_sensor_msgs.hpp>
#include <rviz_common/logging.hpp>
#include <pcl_ros/transforms.hpp>
#include <pcl_conversions/pcl_conversions.h>
#include <tf2_eigen/tf2_eigen.hpp>
#include <pcl/filters/voxel_grid.h>
#include <pcl/common/transforms.h>
#include <QDir>
#include <QProcess>


namespace rviz_custom_plugins
{

namespace utils
{

BagMapLoader::BagMapLoader(QString& path, QObject *parent): QThread(parent), bagPath(path)
{
    qRegisterMetaType<sensor_msgs::msg::PointCloud2::SharedPtr>("sensor_msgs::msg::PointCloud2::SharedPtr");
    qRegisterMetaType<pcl::PointCloud<pcl::PointXYZINormal>::Ptr>("pcl::PointCloud<pcl::PointXYZINormal>::Ptr");
    // qRegisterMetaType<pcl::PointCloud<pcl::Normal>::Ptr>("pcl::PointCloud<pcl::Normal>::Ptr");
    notifier = BagMapNotifier::getInstance();
    connect(this, &BagMapLoader::pointCloudLoaded, notifier, &BagMapNotifier::pointCloudLoaded);
}

void BagMapLoader::run()
{
    // 检查原始bag文件是否存在
    QString bagFilePath = bagPath + "/preview.bag";
    QFileInfo bagFileInfo(bagFilePath);

    {
        std::ostringstream sstm;
        sstm << "bag file path: " << bagFilePath.toStdString();
        RVIZ_COMMON_LOG_INFO(sstm.str().c_str());
    }

    if (!bagFileInfo.exists() || !bagFileInfo.isFile()) {
        return ;
    }
    // 如果bagPath目录下有ros2文件夹，或者文件夹下没有db3文件和yaml文件，则重建文件夹然后运行rosbags-convert命令进行转换
    QDir bagDir(bagPath);
    QString ros2MapPath = bagDir.filePath("ros2");

    {
        std::ostringstream sstm;
        sstm << "ros2 map path: " << ros2MapPath.toStdString();
        RVIZ_COMMON_LOG_INFO(sstm.str().c_str());
    }

    QDir ros2MapDir(ros2MapPath);
    if (!ros2MapDir.exists()) {
        // 构建rosbags-convert命令
        QStringList arguments;
        arguments << "--src" << bagFilePath
                 << "--dst" << ros2MapDir.absolutePath();
        QProcess process;
        process.start("rosbags-convert", arguments);
        if (!process.waitForFinished(5000)) { // 5秒超时

            return ;
        }
        
    }
    
	rosbag2_cpp::Reader reader;
	try {
        reader.open(ros2MapPath.toStdString());
    } catch (const std::exception& e) {
		RVIZ_COMMON_LOG_ERROR("Failed to open map bag");
        return;
    }

	// 指定需要读取的topic
	rosbag2_storage::StorageFilter filter;
	filter.topics = {"/keyframe_pcd", "/keyframe_pose"};
	reader.set_filter(filter);

	std::vector<sensor_msgs::msg::PointCloud2> mapPointCloudVector;
	std::vector<geometry_msgs::msg::PoseStamped> mapPoseVector;

	RVIZ_COMMON_LOG_INFO("read bag file");

	// 逐条读取消息
    while (reader.has_next()) {
        auto bag_msg = reader.read_next();

		if (bag_msg->topic_name == "/keyframe_pcd")  {
			rclcpp::Serialization<sensor_msgs::msg::PointCloud2> serializer;
			rclcpp::SerializedMessage serialized_msg(*bag_msg->serialized_data);
			sensor_msgs::msg::PointCloud2 pcd;
			serializer.deserialize_message(&serialized_msg, &pcd);
			mapPointCloudVector.push_back(pcd);
		} else if (bag_msg->topic_name == "/keyframe_pose") {
			rclcpp::Serialization<geometry_msgs::msg::PoseStamped> serializer;
			rclcpp::SerializedMessage serialized_msg(*bag_msg->serialized_data);
			geometry_msgs::msg::PoseStamped pose;
			serializer.deserialize_message(&serialized_msg, &pose);
			mapPoseVector.push_back(pose);
		}
    }

	if (mapPointCloudVector.size() == 0 || mapPoseVector.size() != mapPointCloudVector.size()) {
        RVIZ_COMMON_LOG_ERROR("pcd size not equal with pose size");
		return;
	}

	RVIZ_COMMON_LOG_INFO("tranform point cloud");

    pcl::PointCloud<pcl::PointXYZINormal>::Ptr pcdPointNormal = mergePointCloud(mapPointCloudVector, mapPoseVector);

    Q_EMIT pointCloudLoaded(pcdPointNormal);
    // Q_EMIT loadPointCloudFinished(std::make_shared<sensor_msgs::msg::PointCloud2>(rosPcd), pcdNormal);
}


// merge逻辑：1.原始点云帧进行体素化并计算法线 2.把所有点云帧合并到一起
// pcl::PointCloud<pcl::PointXYZINormal>::Ptr BagMapLoader::mergePointCloud(std::vector<sensor_msgs::msg::PointCloud2>& mapPointCloudVector,
// 	std::vector<geometry_msgs::msg::PoseStamped>& mapPoseVector)
// {
//     pcl::PointCloud<pcl::PointXYZI>::Ptr pcdPoint(new pcl::PointCloud<pcl::PointXYZI>);
//     pcl::PointCloud<pcl::Normal>::Ptr pclNormal(new pcl::PointCloud<pcl::Normal>);
    
//     for (size_t i = 0; i < mapPointCloudVector.size(); ++i) {
//         // 加入视点信息，把点云转换到world坐标系
// 		auto pcdTrans = transformPointCloudV2(mapPointCloudVector[i], mapPoseVector[i]);
//         auto pcdVoxel = voxelizePcd(pcdTrans, 0.1);
//         *pcdPoint += *pcdVoxel;

//         pcl::PointXYZ viewpoint(mapPoseVector[i].pose.position.x, mapPoseVector[i].pose.position.y, mapPoseVector[i].pose.position.z);
//         auto normals = estimateNormalsFromViewpoint(pcdVoxel, viewpoint);
//         *pclNormal += *normals;
// 	}

//     pcl::PointCloud<pcl::PointXYZINormal>::Ptr pcdPointNormal(new pcl::PointCloud<pcl::PointXYZINormal>);
//     pcl::concatenateFields(*pcdPoint, *pclNormal, *pcdPointNormal);

//     return pcdPointNormal;
// }


// merge逻辑: 1.计算原始点云体素化后的点云 2.遍历体素化后的点云，找到距离最近的视点，根据该视点计算法线
// pcl::PointCloud<pcl::PointXYZINormal>::Ptr BagMapLoader::mergePointCloud(std::vector<sensor_msgs::msg::PointCloud2>& mapPointCloudVector,
// 	std::vector<geometry_msgs::msg::PoseStamped>& mapPoseVector)
// {
//     pcl::PointCloud<pcl::PointXYZI>::Ptr pcdPoint(new pcl::PointCloud<pcl::PointXYZI>);

// 	for (size_t i = 0; i < mapPointCloudVector.size(); ++i) {
//         // 加入视点信息，把点云转换到world坐标系
// 		auto transPcd = transformPointCloudV2(mapPointCloudVector[i], mapPoseVector[i]);
// 		// 追加点云
//         *pcdPoint += *transPcd;
// 	}
//     auto pcdVoxel = voxelizePcd(pcdPoint, 0.1);


//     // 获取飞机相机位置点云
//     pcl::PointCloud<pcl::PointXYZ>::Ptr viewpoints(new pcl::PointCloud<pcl::PointXYZ>);
//     viewpoints->points.reserve(mapPoseVector.size());
//     viewpoints->width = mapPoseVector.size();
//     viewpoints->height = 1;
//     for(auto &pose : mapPoseVector) {
//         pcl::PointXYZ pt;
//         pt.x = pose.pose.position.x;
//         pt.y = pose.pose.position.y;
//         pt.z = pose.pose.position.z;
//         viewpoints->points.push_back(pt);
//     }

//     auto pcdNormal = estimateNormalsFromViewpoints(pcdVoxel, viewpoints);

//     {
//         std::ostringstream sstm;
//         sstm << "normals: " << pcdNormal->size();
//         RVIZ_COMMON_LOG_WARNING(sstm.str().c_str());
//     }

// 	sensor_msgs::msg::PointCloud2 rosPcd;
//     pcl::toROSMsg(*pcdVoxel, rosPcd);
//     // pcl::toROSMsg(*pclPcd, rosPcd);
//     // rosPcd.header = mapPointCloudVector.back().header;

//     pcl::PointCloud<pcl::PointXYZINormal>::Ptr pcdPointNormal(new pcl::PointCloud<pcl::PointXYZINormal>);
//     pcl::concatenateFields(*pcdVoxel, *pcdNormal, *pcdPointNormal);

//     return pcdPointNormal;
// }


// merge逻辑: 1.原始点云计算法线然后合并 2.体素化合并后的原始点云 3.遍历体素化后点云，找到原点云中距离最近的点的法线作为体素化后的法线
// pcl::PointCloud<pcl::PointXYZINormal>::Ptr BagMapLoader::mergePointCloud(std::vector<sensor_msgs::msg::PointCloud2>& mapPointCloudVector,
// 	std::vector<geometry_msgs::msg::PoseStamped>& mapPoseVector)
// {
//     // 原始点云和原始法线
//     pcl::PointCloud<pcl::PointXYZI>::Ptr pcdPoint(new pcl::PointCloud<pcl::PointXYZI>);
//     pcl::PointCloud<pcl::Normal>::Ptr pclNormal(new pcl::PointCloud<pcl::Normal>);
    
//     for (size_t i = 0; i < mapPointCloudVector.size(); ++i) {
//         // 加入视点信息，把点云转换到world坐标系
// 		auto pcdTrans = transformPointCloudV2(mapPointCloudVector[i], mapPoseVector[i]);
//         *pcdPoint += *pcdTrans;

//         pcl::PointXYZ viewpoint(mapPoseVector[i].pose.position.x, mapPoseVector[i].pose.position.y, mapPoseVector[i].pose.position.z);
//         auto normals = estimateNormalsFromViewpoint(pcdTrans, viewpoint);
//         *pclNormal += *normals;
// 	}
//     // 体素化后的点云
//     auto pcdVoxel = voxelizePcd(pcdPoint, 0.1);
//     pcl::PointCloud<pcl::Normal>::Ptr pclVoxelNormal(new pcl::PointCloud<pcl::Normal>);


//     // 创建KDTree对象并设置输入点云B
//     pcl::KdTreeFLANN<pcl::PointXYZI> kdtree;
//     kdtree.setInputCloud(pcdPoint);
    
//     // 用于存储搜索结果的变量
//     std::vector<int> pointIdxNKNSearch(1);      // 最近邻点的索引
//     std::vector<float> pointNKNSquaredDistance(1); // 最近邻点的平方距离
    
//     // 遍历点云A中的每个点
//     for (size_t i = 0; i < pcdVoxel->size(); ++i) {
//         // 执行K近邻搜索，K=1表示只找最近的一个点
//         if (kdtree.nearestKSearch(pcdVoxel->points[i], 1, pointIdxNKNSearch, pointNKNSquaredDistance) > 0) {
//             // 输出结果：点云A中的点i在点云B中的最近邻点是pointIdxNKNSearch[0]
//             // 它们之间的平方距离是pointNKNSquaredDistance[0]
//             int idx = pointNKNSquaredDistance[0];
//             pclVoxelNormal->push_back(pclNormal->normals[idx]);
//         }
//     }

//     pcl::PointCloud<pcl::PointXYZINormal>::Ptr pcdPointNormal(new pcl::PointCloud<pcl::PointXYZINormal>);
//     pcl::concatenateFields(*pcdVoxel, *pclVoxelNormal, *pcdPointNormal);

//     return pcdPointNormal;
// }


// merge逻辑：1.原始点云计算法线然后合并 2.带着法线属性直接体素化 
pcl::PointCloud<pcl::PointXYZINormal>::Ptr BagMapLoader::mergePointCloud(std::vector<sensor_msgs::msg::PointCloud2>& mapPointCloudVector,
	std::vector<geometry_msgs::msg::PoseStamped>& mapPoseVector)
{
    // 原始点云和原始法线
    pcl::PointCloud<pcl::PointXYZI>::Ptr pcdPoint(new pcl::PointCloud<pcl::PointXYZI>);
    pcl::PointCloud<pcl::Normal>::Ptr pclNormal(new pcl::PointCloud<pcl::Normal>);
    
    for (size_t i = 0; i < mapPointCloudVector.size(); ++i) {
        // 加入视点信息，把点云转换到world坐标系
		auto pcdTrans = transformPointCloudV2(mapPointCloudVector[i], mapPoseVector[i]);
        *pcdPoint += *pcdTrans;

        pcl::PointXYZ viewpoint(mapPoseVector[i].pose.position.x, mapPoseVector[i].pose.position.y, mapPoseVector[i].pose.position.z);
        auto normals = estimateNormalsFromViewpoint(pcdTrans, viewpoint);
        *pclNormal += *normals;
	}

    pcl::PointCloud<pcl::PointXYZINormal>::Ptr pcdPointNormal(new pcl::PointCloud<pcl::PointXYZINormal>);
    pcl::concatenateFields(*pcdPoint, *pclNormal, *pcdPointNormal);

    // 体素化后的点云
    pcl::PointCloud<pcl::PointXYZINormal>::Ptr pclVoxelPointNormal = voxelizePcd(pcdPointNormal, 0.1);
    return pclVoxelPointNormal;
}



sensor_msgs::msg::PointCloud2 BagMapLoader::transformPointCloud(
    sensor_msgs::msg::PointCloud2& input_cloud,
    geometry_msgs::msg::PoseStamped& poseStamped)
{
	input_cloud.header.frame_id = "local_odom";
    // 创建 TransformStamped 对象
    geometry_msgs::msg::TransformStamped transform_stamped;
    transform_stamped.header.stamp = input_cloud.header.stamp;
    transform_stamped.header.frame_id = poseStamped.header.frame_id;
    transform_stamped.child_frame_id = input_cloud.header.frame_id;

    // 将 Pose 转换为 Transform
    transform_stamped.transform.translation.x = poseStamped.pose.position.x;
    transform_stamped.transform.translation.y = poseStamped.pose.position.y;
    transform_stamped.transform.translation.z = poseStamped.pose.position.z;
    transform_stamped.transform.rotation = poseStamped.pose.orientation;

    // 使用 tf2 进行点云转换
    sensor_msgs::msg::PointCloud2 output_cloud;
    try {
        // tf2方法性能非常差，请使用pcl库方法
        tf2::doTransform(input_cloud, output_cloud, transform_stamped);
    } catch (const tf2::TransformException &ex) {
        RCLCPP_ERROR(rclcpp::get_logger("transform_pointcloud"),
                    "Failed to transform point cloud: %s", ex.what());
        return input_cloud;
    }

    return output_cloud;
}

pcl::PointCloud<pcl::PointXYZI>::Ptr BagMapLoader::transformPointCloudV2(
	sensor_msgs::msg::PointCloud2& inputCloud,
	geometry_msgs::msg::PoseStamped& poseStamped)
{

	tf2::Quaternion q(poseStamped.pose.orientation.x,
		poseStamped.pose.orientation.y,
		poseStamped.pose.orientation.z,
		poseStamped.pose.orientation.w);
	// tf2::Matrix3x3 rot_mat_tf(q);
	Eigen::Quaterniond quat_eigen(q.w(), q.x(), q.y(), q.z());
	Eigen::Matrix3d rot_mat_eig = quat_eigen.toRotationMatrix();
	// tf2::matrixTFToEigen(rot_mat_tf, rot_mat_eig);
	Eigen::Matrix4d pose_eig_ = Eigen::Matrix4d::Identity();
	pose_eig_.block<3, 3>(0, 0) = rot_mat_eig;
	pose_eig_(0, 3) = poseStamped.pose.position.x;
	pose_eig_(1, 3) = poseStamped.pose.position.y;
	pose_eig_(2, 3) = poseStamped.pose.position.z;
	pcl::PointCloud<pcl::PointXYZI> pcd;
	pcl::fromROSMsg(inputCloud, pcd);

    pcl::PointCloud<pcl::PointXYZI>::Ptr outputPcd(new pcl::PointCloud<pcl::PointXYZI>);
    outputPcd->reserve(pcd.size());
	// pcl::PointCloud<pcl::PointXYZI> outputPcd = pcd;
    pcl::transformPointCloud(pcd, *outputPcd, pose_eig_);
	return outputPcd;
}


pcl::PointCloud<pcl::Normal>::Ptr BagMapLoader::estimateNormalsFromViewpoints(
    pcl::PointCloud<pcl::PointXYZI>::Ptr cloud,
    pcl::PointCloud<pcl::PointXYZ>::Ptr viewpoints,
    int kSearch)
{
	// 创建法线估计对象
	pcl::NormalEstimation<pcl::PointXYZI, pcl::Normal> ne;
	ne.setInputCloud(cloud);
	// 设置搜索方法
	pcl::search::KdTree<pcl::PointXYZI>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZI>());
	ne.setSearchMethod(tree);
	// 设置邻域点数
	ne.setKSearch(kSearch);
    // 计算法线
    pcl::PointCloud<pcl::Normal>::Ptr normals(new pcl::PointCloud<pcl::Normal>);
	ne.compute(*normals);

    pcl::KdTreeFLANN<pcl::PointXYZ> viewKdTree;
    viewKdTree.setInputCloud(viewpoints);
    // 修正法线方向，使其朝向离自己最近的视点
    for (size_t i = 0; i < cloud->size(); ++i) {
        const pcl::PointXYZI& pt = cloud->points[i];
        pcl::Normal& n = normals->points[i];

        pcl::PointXYZ searchPt(pt.x, pt.y, pt.z);
        std::vector<int> indices(1);
        std::vector<float> dists(1);
        viewKdTree.nearestKSearch(searchPt, 1, indices, dists);

        Eigen::Vector3f viewpoint(
            viewpoints->points[indices[0]].x,
            viewpoints->points[indices[0]].y,
            viewpoints->points[indices[0]].z
        );

        Eigen::Vector3f p(pt.x, pt.y, pt.z);
        Eigen::Vector3f normal(n.normal_x, n.normal_y, n.normal_z);
        Eigen::Vector3f to_view = viewpoint - p;

        if (normal.dot(to_view) < 0){
            normal = -normal;
        }

        n.normal_x = normal.x();
        n.normal_y = normal.y();
        n.normal_z = normal.z();
    }
    return normals;
}


pcl::PointCloud<pcl::Normal>::Ptr BagMapLoader::estimateNormalsFromViewpoint(
    pcl::PointCloud<pcl::PointXYZI>::Ptr cloud,
    pcl::PointXYZ& viewpoint,
    int kSearch)
{
    pcl::NormalEstimation<pcl::PointXYZI, pcl::Normal> ne;
    ne.setInputCloud(cloud);

    pcl::search::KdTree<pcl::PointXYZI>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZI>);
    ne.setSearchMethod(tree);
    ne.setKSearch(kSearch); 

    // 设置统一视点（默认是 (0,0,0)，你可以指定相机位置）
    ne.setViewPoint(viewpoint.x, viewpoint.y, viewpoint.z);

    // 输出点云
    pcl::PointCloud<pcl::Normal>::Ptr normals(new pcl::PointCloud<pcl::Normal>);
    ne.compute(*normals);
    return normals;
}


// pcl::PointCloud<pcl::PointXYZI>::Ptr BagMapLoader::voxelizePcd(const pcl::PointCloud<pcl::PointXYZI>::Ptr pcdIn, const float voxelRes)
// {
//     pcl::VoxelGrid<pcl::PointXYZI> voxelGrid;
//     voxelGrid.setLeafSize(voxelRes, voxelRes, voxelRes);
//     // pcl::PointCloud<pcl::PointXYZI>::Ptr pcd_in_ptr(new pcl::PointCloud<pcl::PointXYZI>);
//     pcl::PointCloud<pcl::PointXYZI>::Ptr pcdOut(new pcl::PointCloud<pcl::PointXYZI>);
//     // pcd_in_ptr->reserve(pcd_in.size());
//     pcdOut->reserve(pcdIn->size());
//     // *pcd_in_ptr = pcd_in;
//     voxelGrid.setInputCloud(pcdIn);
//     voxelGrid.filter(*pcdOut);
//     return pcdOut;
// }

pcl::PointCloud<pcl::PointXYZINormal>::Ptr BagMapLoader::voxelizePcd(const pcl::PointCloud<pcl::PointXYZINormal>::Ptr pcdIn, const float voxelRes)
{
    pcl::VoxelGrid<pcl::PointXYZINormal> voxelGrid;
    pcl::PointCloud<pcl::PointXYZINormal>::Ptr pcdOut(new pcl::PointCloud<pcl::PointXYZINormal>);
    pcdOut->reserve(pcdIn->size());
    // *pcd_in_ptr = pcd_in;
    voxelGrid.setInputCloud(pcdIn);
    voxelGrid.setLeafSize(voxelRes, voxelRes, voxelRes);
    voxelGrid.setDownsampleAllData(true);
    voxelGrid.filter(*pcdOut);
    return pcdOut;
}

}


}
